# 导入所需的库
import argparse  # 用于解析命令行参数
import os  # 处理文件和目录路径
import time  # 处理时间相关功能
import json  # 用于保存参数为JSON格式
from pathlib import Path  # 面向对象的文件系统路径操作

import gymnasium as gym  # 强化学习环境接口
from stable_baselines3 import PPO  # PPO算法实现
from stable_baselines3.common.callbacks import EvalCallback  # 评估回调函数
from stable_baselines3.common.vec_env import SubprocVecEnv  # 多进程并行环境
from stable_baselines3.common.env_util import make_vec_env  # 创建并行环境工具函数

# from go1_mujoco_env import Go1MujocoEnv  # 自定义的Go1机器人Mujoco环境
from bittle_mujoco_env import BittleMujocoEnv  # 自定义的Bittle机器人Mujoco环境  #XXX
from tqdm import tqdm  # 显示进度条

# 定义模型和日志保存目录
MODEL_DIR = "models"
LOG_DIR = "logs"


def train(args):
    """训练PPO模型"""
    # 创建并行环境，使用SubprocVecEnv实现多进程并行
    vec_env = make_vec_env(
        # Go1MujocoEnv,  # 环境类
        BittleMujocoEnv,  # 环境类 #XXX
        env_kwargs={"ctrl_type": args.ctrl_type},  # 传递给环境的参数（控制类型）
        n_envs=args.num_parallel_envs,  # 并行环境数量
        seed=args.seed,  # 随机种子
        vec_env_cls=SubprocVecEnv,  # 并行环境类
    )

    # 生成运行名称，包含时间戳和自定义名称
    train_time = time.strftime("%Y-%m-%d_%H-%M-%S")
    run_name = (
        f"{train_time}" if args.run_name is None else f"{train_time}-{args.run_name}"
    )

    # 模型保存路径
    model_path = f"{MODEL_DIR}/{run_name}"
    os.makedirs(model_path, exist_ok=True)  # 确保目录存在
    print(
        f"Training on {args.num_parallel_envs} parallel environments, saving models to '{model_path}'"
    )
    # 保存训练参数（奖励/成本权重等）
    try:
        # 使用 unwrapped 或 get_wrapper_attr 提取参数
        reward_weights = vec_env.get_attr("reward_weights")[
            0
        ]  # 如果 reward_weights 是环境属性
        cost_weights = vec_env.get_attr("cost_weights")[
            0
        ]  # 如果 cost_weights 是环境属性
    except Exception as e:
        print(f"Warning: Failed to get weights from env - {e}")
        reward_weights = "Not Available"
        cost_weights = "Not Available"

    # 构建参数字典
    params = {
        "reward_weights": reward_weights,
        "cost_weights": cost_weights,
        "ctrl_type": args.ctrl_type,
        "seed": args.seed,
        "num_parallel_envs": args.num_parallel_envs,
    }

    # 保存参数为JSON文件
    params_path = os.path.join(model_path, "training_params.json")
    with open(params_path, "w") as f:
        json.dump(params, f, indent=4)
    print(f"Training parameters saved to {params_path}")

    # 创建评估回调函数，定期评估模型性能
    eval_callback = EvalCallback(
        vec_env,
        best_model_save_path=model_path,  # 最佳模型保存路径
        log_path=LOG_DIR,  # 评估日志路径
        eval_freq=args.eval_frequency,  # 每训练多少步评估一次
        n_eval_episodes=5,  # 每次评估运行5个episode
        deterministic=True,  # 使用确定性动作
        render=True,  # 不渲染评估过程
    )

    # 加载现有模型或创建新模型
    if args.model_path:
        model = PPO.load(
            args.model_path, env=vec_env, verbose=1, tensorboard_log=LOG_DIR
        )
    else:
        # 创建新的PPO模型，使用MLP策略
        model = PPO(
            "MlpPolicy",
            vec_env,
            verbose=1,  # 输出训练信息
            tensorboard_log=LOG_DIR,  # TensorBoard日志目录
        )

    # 开始训练
    model.learn(
        total_timesteps=args.total_timesteps,  # 总训练步数
        reset_num_timesteps=False,  # 是否重置步数计数器
        progress_bar=True,  # 显示进度条
        tb_log_name=run_name,  # TensorBoard运行名称
        callback=eval_callback,  # 使用评估回调
    )
    # 保存最终模型
    model.save(f"{model_path}/final_model")


def test(args):
    """测试训练好的模型"""

    # 如果没有指定模型路径，自动选择最新的模型路径
    if not args.model_path:
        # 获取models目录下所有文件夹（模型文件夹）
        model_dir = Path(MODEL_DIR)
        model_folders = [f for f in model_dir.iterdir() if f.is_dir()]

        if not model_folders:
            raise ValueError("No model found in MODEL_DIR.Train one first.")

        # 按照创建时间排序，选择最新的模型文件夹
        latest_model_folder = max(model_folders, key=lambda f: f.stat().st_mtime)
        model_path = latest_model_folder / "final_model.zip"  # 默认模型文件名

        print(f"Using latest model for test: {model_path}")
    else:
        model_path = Path(args.model_path)

    # 配置测试环境（实时渲染或视频录制）
    if not args.record_test_episodes:
        # 实时渲染模式
        # env = Go1MujocoEnv(
        #     ctrl_type=args.ctrl_type,
        #     render_mode="human",  # 人类可实时查看的渲染模式
        # )

        env = BittleMujocoEnv(  # XXX
            ctrl_type=args.ctrl_type,
            render_mode="human",  # 人类可实时查看的渲染模式
        )

        inter_frame_sleep = 0.016  # 模拟实时渲染的帧间隔（约60FPS）
    else:
        # 视频录制模式
        # env = Go1MujocoEnv(
        #     ctrl_type=args.ctrl_type,
        #     render_mode="rgb_array",  # 返回RGB数组用于录制
        #     camera_name="tracking",  # 使用的相机名称
        #     width=1920,  # 视频宽度
        #     height=1080,  # 视频高度
        # )

        env = BittleMujocoEnv(  # XXX
            ctrl_type=args.ctrl_type,
            render_mode="rgb_array",  # 返回RGB数组用于录制
            camera_name="tracking",  # 使用的相机名称
            width=1920,  # 视频宽度
            height=1080,  # 视频高度
        )

        # 添加视频录制包装器
        env = gym.wrappers.RecordVideo(
            env,
            video_folder="recordings/",  # 视频保存目录
            name_prefix=model_path.parent.name,  # 视频文件名前缀
        )
        inter_frame_sleep = 0.0  # 无需延迟

    # 加载训练好的模型
    model = PPO.load(path=model_path, env=env, verbose=1)

    # 进行指定次数的测试
    num_episodes = args.num_test_episodes
    total_reward = 0
    total_length = 0
    for _ in tqdm(range(num_episodes)):
        obs, _ = env.reset()
        env.render()  # 初始化渲染

        ep_len = 0
        ep_reward = 0
        while True:
            action, _ = model.predict(obs, deterministic=True)  # 使用确定性策略
            obs, reward, terminated, truncated, info = env.step(action)
            ep_reward += reward
            ep_len += 1

            time.sleep(inter_frame_sleep)  # 控制渲染速度

            if terminated or truncated:  # episode结束条件
                print(f"{ep_len=}  {ep_reward=}")
                break

        # 累计统计信息
        total_length += ep_len
        total_reward += ep_reward

    # 输出平均性能指标
    print(
        f"Avg reward: {total_reward/num_episodes}, Avg length: {total_length/num_episodes}"
    )


if __name__ == "__main__":
    # 配置命令行参数解析器
    # parser = argparse.ArgumentParser(
    #     description="Train or test PPO on Go1 Mujoco environment."
    # )

    parser = argparse.ArgumentParser(
        description="Train or test PPO on Bittle Mujoco environment."  # XXX
    )

    parser.add_argument(
        "--run",
        type=str,
        required=True,
        choices=["train", "test"],
        help="Select training or testing mode",
    )

    # Training related parameters
    parser.add_argument(
        "--run_name",
        type=str,
        default=None,
        help="The name of the training run, used for saving models and logs",
    )
    parser.add_argument(
        "--num_parallel_envs",
        type=int,
        default=8,
        help="Number of parallel environments during training",
    )
    parser.add_argument(
        "--total_timesteps",
        type=int,
        default=5_000_000,
        help="Total number of training timesteps",
    )
    parser.add_argument(
        "--eval_frequency",
        type=int,
        default=100_000,
        help="Frequency (in timesteps) for evaluating the model",
    )

    # Testing related parameters
    parser.add_argument(
        "--num_test_episodes",
        type=int,
        default=5,
        help="Number of episodes to run during testing",
    )
    parser.add_argument(
        "--record_test_episodes",
        action="store_true",
        help="Whether to record test episodes as videos",
    )

    # General parameters
    parser.add_argument(
        "--model_path",
        type=str,
        default=None,
        help="Path to the model (for continuing training or testing)",
    )
    parser.add_argument(
        "--ctrl_type",
        choices=["torque", "position"],
        default="position",
        help="Robot control type: torque control or position control",
    )
    parser.add_argument("--seed", type=int, default=0, help="Random seed")

    args = parser.parse_args()

    # 根据参数选择运行模式
    if args.run == "train":
        os.makedirs(MODEL_DIR, exist_ok=True)
        os.makedirs(LOG_DIR, exist_ok=True)
        train(args)
    elif args.run == "test":
        # if not args.model_path:
        # raise ValueError("测试需要指定--model_path参数")
        test(args)
